import hmmtop
import unittest
from Bio import SeqIO

class TestHmmtop(unittest.TestCase):

  def test_seq1(self):
    record = SeqIO.read(open("P55117.fasta"), "fasta")
    hmmtop_res = hmmtop.hmmtop('>' + str(record.id), str(record.seq))
    self.assertEqual(hmmtop_res.success, True)
    self.assertEqual(hmmtop_res.num_tms, 2)
    self.assertEqual(hmmtop_res.tms_ranges[0], (298, 320))
    self.assertEqual(hmmtop_res.tms_ranges[1], (378, 403))
    self.assertEqual(hmmtop_res.tms_seq[0], 'VAALIASTVALAISPLSFAGIA')
    self.assertEqual(hmmtop_res.tms_seq[1], 'AAGSVVASPIALLVSGITGVISTIL')
    self.assertEqual(hmmtop_res.tms_seq[0], 'VAALIASTVALAISPLSFAGIA')

    is_tms = hmmtop_res.get_is_tms('T--GP-V-AALIASTVALAI-SPLS----')
    self.assertEqual(is_tms, [False, '-', '-', False, False, '-', True, '-', True, True, True, True, True, True, True, True, True, True, True, True, '-', True, True, True, True, '-', '-', '-', '-'])
    self.assertEqual(hmmtop_res.get_is_tms('not even a peptide'), False)


  def test_seq2(self):
    record = SeqIO.read(open("Q9S6V3.fasta"), "fasta")
    hmmtop_res = hmmtop.hmmtop(str(record.id), str(record.seq))
    self.assertEqual(hmmtop_res.success, True)
    self.assertEqual(hmmtop_res.num_tms, 12)
    self.assertEqual(hmmtop_res.tms_ranges[0], (26, 44))
    self.assertEqual(hmmtop_res.tms_ranges[1], (62, 82))
    self.assertEqual(hmmtop_res.tms_ranges[2], (94, 114))
    self.assertEqual(hmmtop_res.tms_ranges[3], (146, 171))
    self.assertEqual(hmmtop_res.tms_ranges[4], (183, 204))
    self.assertEqual(hmmtop_res.tms_ranges[5], (222, 242))
    self.assertEqual(hmmtop_res.tms_ranges[6], (283, 302))
    self.assertEqual(hmmtop_res.tms_ranges[7], (320, 343))
    self.assertEqual(hmmtop_res.tms_ranges[8], (355, 375))
    self.assertEqual(hmmtop_res.tms_ranges[9], (379, 399))
    self.assertEqual(hmmtop_res.tms_ranges[10], (446, 464))
    self.assertEqual(hmmtop_res.tms_ranges[11], (468, 486))

    self.assertEqual(hmmtop_res.tms_seq[0], 'VLPMFLMFFCISFNYTIL')
    self.assertEqual(hmmtop_res.tms_seq[1], 'IPFIKLWLVVPSAVVFMLIY')
    self.assertEqual(hmmtop_res.tms_seq[2], 'FFAVLSPFVVFFALFPVVIY')
    self.assertEqual(hmmtop_res.tms_seq[3], 'WTFAVFYVLSELWGSVMLSLMFWGF')
    self.assertEqual(hmmtop_res.tms_seq[4], 'FYALFGVGANVALLISGPAII')
    self.assertEqual(hmmtop_res.tms_seq[5], 'SLYFLMAMFLCSCAIIAACY')
    self.assertEqual(hmmtop_res.tms_seq[6], 'YMLLLALLVICYGICINLV')
    self.assertEqual(hmmtop_res.tms_seq[7], 'SAFMGNFSFWTGVVSVFVMLFIG')
    self.assertEqual(hmmtop_res.tms_seq[8], 'ALVTPIMVLVTGAVFFALVI')
    self.assertEqual(hmmtop_res.tms_seq[9], 'ATGLVAALGTTPLMLAVVVG')
    self.assertEqual(hmmtop_res.tms_seq[10], 'SGGSLIQQGLLVVCGSIG')
    self.assertEqual(hmmtop_res.tms_seq[11], 'FLAVALFAIIMVWLTSAT')

    
    is_tms = hmmtop_res.get_is_tms('--LFI--GG-NVI-R--')
    self.assertEqual(is_tms, ['-', '-', True, True, True, '-', '-', True, False, '-', False, False, False, '-', False, '-', '-'])

if __name__ == '__main__':
  unittest.main()
